import torch
import scipy
import numpy as np

from torch.utils.data import DataLoader

from .TrainUtil import *
from .DataProcess import *
from .ModelUtil import *

import os

def Label_attack(dataloader, target_model, loss_fn, device, sigma_list, nums):    
    result_list = []
    target_result = noise_infer(dataloader, target_model, loss_fn, device, sigma=0)
    for sigma in sigma_list:
        for i in range(nums):
            pred_result = noise_infer(dataloader, target_model, loss_fn, device, sigma=sigma)
            pred_result = pred_result.detach().cpu().numpy()
            result_list.append(pred_result)

    result_list = np.array(result_list)
    result_sum = np.sum(result_list, axis=0)
    target_result = target_result.detach().cpu().numpy()
    noise_result = np.multiply(target_result, result_sum)
    noise_result = noise_result/(nums*len(sigma_list))
    pred_result = (noise_result>=np.median(noise_result))
    return pred_result, noise_result


# 基线攻击
def base_attack(dataloader, target_model, loss_fn, device):    
    pred_result = noise_infer(dataloader, target_model, loss_fn, device, sigma=0)
    pred_result = pred_result.detach().cpu().numpy()
    return pred_result